#!/usr/bin/env python3
import os
from fastapi import FastAPI, File, UploadFile, Request, HTTPException, Form
from fastapi.responses import HTMLResponse, FileResponse, JSONResponse
from fastapi.templating import Jinja2Templates
from datetime import datetime
import math
import argparse

parser = argparse.ArgumentParser(description='http简易文件服务器')
parser.add_argument('-p', '--port', type=int, default=18181, help='指定启动的端口，默认:18181')
parser.add_argument('-M', '--max-file-size', type=int, default=20, help='修改可上传的最大文件大小,单位:GB,默认:20')
parser.add_argument('-d', '--directory', default='uploads', help='指定要上传的目录名，默认:uploads')
parser.add_argument('-c', '--chunk_size', type=int, default=5,
                    help='分块上传，每块儿的大小，大文件可取大一些，单位:MB，默认值:5')

args = parser.parse_args()

# 配置
UPLOAD_DIR = args.directory  # 上传文件目录
CHUNK_SIZE = args.chunk_size * 1024 * 1024  # 分块大小，5MB (与前端保持一致)
MAX_FILE_SIZE = args.max_file_size * 1024 * 1024 * 1024  # 最大文件大小，20GB

# 创建上传目录
if not os.path.exists(UPLOAD_DIR):
    os.makedirs(UPLOAD_DIR)

app = FastAPI(title="文件上传服务")

# 创建Jinja2环境并注册过滤器
templates = Jinja2Templates(directory="templates")

# 文件类型图标映射
FILE_ICON_MAP = {
    'pdf': 'fa-file-pdf-o',
    'doc': 'fa-file-word-o',
    'docx': 'fa-file-word-o',
    'xls': 'fa-file-excel-o',
    'xlsx': 'fa-file-excel-o',
    'ppt': 'fa-file-powerpoint-o',
    'pptx': 'fa-file-powerpoint-o',
    'jpg': 'fa-file-image-o',
    'jpeg': 'fa-file-image-o',
    'png': 'fa-file-image-o',
    'gif': 'fa-file-image-o',
    'bmp': 'fa-file-image-o',
    'webp': 'fa-file-image-o',
    'zip': 'fa-file-archive-o',
    'rar': 'fa-file-archive-o',
    '7z': 'fa-file-archive-o',
    'txt': 'fa-file-text-o',
}


# 注册日期格式化过滤器
def datetimeformat(value, format='%Y-%m-%d %H:%M:%S'):
    return datetime.fromtimestamp(value).strftime(format)


# 获取文件图标类
def get_file_icon_class(filename: str) -> str:
    ext = filename.split('.')[-1].lower() if '.' in filename else ''
    return FILE_ICON_MAP.get(ext, 'fa-file-o')


# 注册过滤器和全局函数
templates.env.filters['datetimeformat'] = datetimeformat
templates.env.globals['get_file_icon_class'] = get_file_icon_class

# 存储文件上传状态
file_upload_status = {}


@app.get("/", response_class=HTMLResponse)
async def index(request: Request):
    """返回上传页面"""
    # 获取已上传文件列表
    files = []
    if os.path.exists(UPLOAD_DIR):
        for filename in os.listdir(UPLOAD_DIR):
            file_path = os.path.join(UPLOAD_DIR, filename)
            if os.path.isfile(file_path) and not filename.endswith(".part"):
                files.append({
                    "name": filename,
                    "size": format_size(os.path.getsize(file_path)),
                    "mtime": os.path.getmtime(file_path)
                })

    # 按修改时间排序，最新的在前
    files.sort(key=lambda x: x["mtime"], reverse=True)

    return templates.TemplateResponse("index.html", {
        "request": request,
        "files": files,
        "chunk_size": args.chunk_size,
        "max_file_size": format_size(MAX_FILE_SIZE)
    })


@app.post("/upload")
def upload_file(
        request: Request,
        file: UploadFile = File(...),
        file_name: str = Form(...),
        chunk_index: int = Form(...),
        total_chunks: int = Form(...),
        action: str = Form("new")  # new: 新上传, overwrite: 覆盖, resume: 续传
):
    """处理文件上传（支持分块）"""
    try:

        # 验证文件大小
        if action == "new" or action == "overwrite":
            if file.size * total_chunks > MAX_FILE_SIZE:
                return JSONResponse(
                    {"error": f"文件大小超过限制 ({format_size(MAX_FILE_SIZE)})"},
                    status_code=413
                )

        file_path = os.path.join(UPLOAD_DIR, file_name)

        # 处理新上传或覆盖 - 重置状态
        if action == "new" and chunk_index == 0:
            # 重置上传状态
            file_upload_status[file_name] = {
                "total_chunks": total_chunks,
                "received_chunks": 0,
                "file_path": file_path
            }
        if action == "overwrite" and chunk_index == 0:
            if os.path.exists(file_path):
                os.remove(file_path)
            # 重置上传状态
            file_upload_status[file_name] = {
                "total_chunks": total_chunks,
                "received_chunks": 0,
                "file_path": file_path
            }

        # 处理续传操作
        if action == "resume":
            if file_name not in file_upload_status:
                # 尝试从磁盘恢复状态
                if os.path.exists(file_path):
                    file_size = os.path.getsize(file_path)
                    received_chunks = math.ceil(file_size / CHUNK_SIZE)
                    file_upload_status[file_name] = {
                        "total_chunks": total_chunks,
                        "received_chunks": received_chunks,
                        "file_path": file_path
                    }
                else:
                    return {"status": "error", "message": "无法恢复上传状态"}

        # 确保状态存在
        if file_name not in file_upload_status:
            return {"status": "error", "message": "上传状态不存在"}

        # 获取当前状态
        status = file_upload_status[file_name]

        # 验证分块索引
        if chunk_index != status["received_chunks"]:
            return {"status": "error",
                    "message": f"错误的分块索引, 期望 {status['received_chunks']}, 收到 {chunk_index}"}

        # 读取文件内容（同步）
        chunk_data = file.file.read()

        # 追加写入分块内容
        with open(status["file_path"], "ab") as f:
            f.write(chunk_data)

        # 更新状态
        status["received_chunks"] += 1
        received_chunks = status["received_chunks"]

        # 检查是否所有分块都已接收
        if received_chunks == total_chunks:
            # 清理状态
            del file_upload_status[file_name]

            return {"status": "success", "message": "文件上传完成", "filename": file_name}

        return {
            "status": "success",
            "message": f"分块 {received_chunks}/{total_chunks} 上传成功",
            "next_chunk": received_chunks
        }

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"上传失败: {str(e)}")


@app.get("/get_resume_info")
async def get_resume_info(file_name: str):
    """获取文件续传信息"""
    file_path = os.path.join(UPLOAD_DIR, file_name)

    # 检查文件是否存在
    file_exists = os.path.exists(file_path)

    # 检查是否有部分上传的文件
    uploaded_bytes = 0
    uploaded_chunks = 0

    if os.path.exists(file_path):
        uploaded_bytes = os.path.getsize(file_path)
        uploaded_chunks = math.ceil(uploaded_bytes / CHUNK_SIZE)

    return {
        "file_name": file_name,
        "file_exists": file_exists,
        "uploaded_bytes": uploaded_bytes,
        "uploaded_chunks": uploaded_chunks
    }


@app.get("/download/{filename}")
async def download_file(filename: str):
    """下载文件"""
    file_path = os.path.join(UPLOAD_DIR, filename)
    if os.path.isfile(file_path):
        return FileResponse(file_path, filename=filename)
    raise HTTPException(status_code=404, detail="文件不存在")


@app.delete("/delete/{filename}")
async def delete_file(filename: str):
    """删除文件"""
    file_path = os.path.join(UPLOAD_DIR, filename)
    temp_file_path = file_path + ".part"

    # 删除主文件
    if os.path.isfile(file_path):
        os.remove(file_path)

    # 删除临时文件
    if os.path.isfile(temp_file_path):
        os.remove(temp_file_path)

    # 删除上传状态
    if filename in file_upload_status:
        del file_upload_status[filename]

    return {"status": "success", "message": "文件已删除"}


@app.get("/static/{path:path}")
async def static_file(path: str):
    """获取静态文件（支持多级目录）"""
    file_path = os.path.join("static", path)
    if os.path.isfile(file_path):
        return FileResponse(file_path)
    raise HTTPException(status_code=404, detail="文件不存在")


# 工具函数：格式化文件大小
def format_size(size_bytes):
    """将字节转换为可读的文件大小格式"""
    units = ["B", "KB", "MB", "GB", "TB"]
    unit_index = 0
    while size_bytes >= 1024 and unit_index < len(units) - 1:
        size_bytes /= 1024
        unit_index += 1
    return f"{size_bytes:.2f} {units[unit_index]}"


# 启动命令提示
if __name__ == "__main__":
    import uvicorn

    print(f"服务器运行在 http://0.0.0.0:{args.port}")
    print(f"上传的文件将保存在 {os.path.abspath(UPLOAD_DIR)} 目录")
    print(f"最大文件大小: {format_size(MAX_FILE_SIZE)}")
    uvicorn.run(app, host="0.0.0.0", port=args.port)
